# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
import os
import time
import moxing as mox
import mindspore.dataset as ds
from mindspore import context, Model
from mindspore.train.serialization import load_checkpoint, load_param_into_net

from models.net import DAM_Net, PredictWithSigmoid, DAMNetWithLoss, DAMTrainOneStepCell, ModelBuilder
from utils.metric import UbuntuTestMetric, DoubanTestMetric, EvalMetric
import utils.config as config

mox.file.shift('os', 'mox')

device_num = int(os.getenv('RANK_SIZE'))
device_id = int(os.getenv('DEVICE_ID'))
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")


def eval(args):
    config = args
    print(config, '\n')

    root = "/cache/"
    obs_data_path = config.data_url
    obs_ckpt_path = config.checkpoint_path
    obs_test_path = os.path.join(config.train_url, str(config.time_))
    if config.model_name == "DAM_ubuntu":
        local_data_path = os.path.join(root, "ubuntu_data")
    else:
        local_data_path = os.path.join(root, "douban_data")
    local_ckpt_path = os.path.join(local_data_path, "ckpt")
    local_test_path = os.path.join(root, "test")
    mox.file.make_dirs(local_data_path)
    mox.file.make_dirs(local_ckpt_path)
    mox.file.make_dirs(local_test_path)

    print("############## Downloading data from OBS ##############")
    mox.file.copy_parallel(src_url=obs_data_path, dst_url=local_data_path)  # 将obs的数据拷贝到云上的本地
    mox.file.copy_parallel(src_url=obs_ckpt_path, dst_url=local_ckpt_path)
    print("############### Downloading is completed ##############")

    test_data_path = os.path.join(local_data_path, "data_test.mindrecord")
    print("************Starting loading data: ", test_data_path)
    print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())))
    dataset = ds.MindDataset(test_data_path,
                             columns_list=["turns", "turn_len", "response", "response_len", "label"],
                             shuffle=False, num_shards=device_num, shard_id=device_id)
    dataset = dataset.batch(config.eval_batch_size, drop_remainder=True)
    dataset = dataset.repeat(1)
    print("dataset: ", dataset)
    print("dataset_len: ", dataset.get_dataset_size() * config.eval_batch_size)
    print("dataset_size: ", dataset.get_dataset_size())
    print("*************Finish loading data**************")
    print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())))

    dam_net = DAM_Net(config, train_mode=False)
    train_net = DAMNetWithLoss(dam_net)
    train_net = DAMTrainOneStepCell(train_net, config)
    eval_net = PredictWithSigmoid(dam_net)
    if config.model_name == 'DAM_ubuntu':
        metric = EvalMetric("DAM_ubuntu")
    else:
        metric = EvalMetric("DAM_douban")
    model = Model(train_net, eval_network=eval_net, metrics={"Accuracy": metric})

    # 加载checkpoint
    checkpoint_file = os.path.join(local_ckpt_path, config.checkpoint_name)
    print('loading checkpoint: ', checkpoint_file)
    param_dict = load_checkpoint(checkpoint_file)
    load_param_into_net(dam_net, param_dict)

    ckpt_name = config.checkpoint_name
    ckpt_name = ckpt_name.split('.')[0]
    test_file = os.path.join(local_test_path, "result_" + ckpt_name + ".test")
    print("test_file: ", test_file)

    print('testing started...')
    res = model.eval(dataset, dataset_sink_mode=False)
    print(res)
    result = res["Accuracy"]

    with open(test_file, 'a+', encoding='utf-8') as file_out:
        file_out.write("checkpoint_file: " + config.checkpoint_path + config.checkpoint_name + '\n')
        for acc in result:
            file_out.write(str(acc) + '\n')

    mox.file.copy_parallel(src_url=local_test_path, dst_url=obs_test_path)


if __name__ == '__main__':
    # eval(config.douban_parse_args())
    eval(config.ubuntu_parse_args())

